# coding=utf8
from __future__ import division
import os
import os.path
from glob import glob
import sys
import torch
import torch.utils.data as data
import numpy as np
import pandas as pd
from matplotlib import pyplot as plt
from glob import glob
from augment import FaceAug
import cv2
import zipfile
from matlab_cp2tform import get_similarity_transform_for_cv2

def alignment(src_img,src_pts):
    ref_pts = [ [30.2946, 51.6963],[65.5318, 51.5014],
        [48.0252, 71.7366],[33.5493, 92.3655],[62.7299, 92.2041] ]
    crop_size = (96, 112)
    src_pts = np.array(src_pts).reshape(5,2)

    s = np.array(src_pts).astype(np.float32)
    r = np.array(ref_pts).astype(np.float32)

    tfm = get_similarity_transform_for_cv2(s, r)
    face_img = cv2.warpAffine(src_img, tfm, crop_size)
    return face_img

class LFW_np(data.Dataset):
    def __init__(self, images, labels):
        self.images = images  #(N,2,112,96,1)  (N,2,H,W,1)
        self.labels = labels

    def __len__(self):
        return self.images.shape[0]

    def __getitem__(self, item):

        return torch.from_numpy(self.images[item,0,:]).float(), \
               torch.from_numpy(self.images[item, 1, :]).float(),\
               self.labels[item]

class LFW_zip(data.Dataset):
    def __init__(self, lfw_path,pairtxt, landmarktxt):
        with open(pairtxt) as f:
            pairs_lines = f.readlines()[1:]
        self.pairs = pairs_lines
        self.zfile = zipfile.ZipFile(lfw_path)

        landmark = {}
        with open(landmarktxt) as f:
            landmark_lines = f.readlines()
        for line in landmark_lines:
            l = line.replace('\n', '').split('\t')
            landmark[l[0]] = [int(k) for k in l[1:]]

        self.landmark = landmark
    def __len__(self):
        return len(self.pairs)

    def __getitem__(self, item):
        p = self.pairs[item].replace('\n', '').split('\t')
        if 3 == len(p):
            label = 1
            name1 = p[0] + '/' + p[0] + '_' + '{:04}.jpg'.format(int(p[1]))
            name2 = p[0] + '/' + p[0] + '_' + '{:04}.jpg'.format(int(p[2]))
        elif 4 == len(p):
            label = 0
            name1 = p[0] + '/' + p[0] + '_' + '{:04}.jpg'.format(int(p[1]))
            name2 = p[2] + '/' + p[2] + '_' + '{:04}.jpg'.format(int(p[3]))

        else:
            raise ValueError('invalid line:"%s"'%self.pairs[item])

        # (112,96,3)
        img1 = alignment(cv2.imdecode(np.frombuffer(self.zfile.read(name1), np.uint8), 1), self.landmark[name1])
        img2 = alignment(cv2.imdecode(np.frombuffer(self.zfile.read(name2), np.uint8), 1), self.landmark[name2])

        img1 = img1.transpose(2, 0, 1)
        img2 = img2.transpose(2, 0, 1)

        img1 = (img1 - 127.5) / 128.0
        img2 = (img2 - 127.5) / 128.0

        return torch.from_numpy(img1).float(), \
               torch.from_numpy(img2).float(),\
               label

class MS1M(data.Dataset):
    def __init__(self,root_path, list_csv, transform=None):
        listpd = pd.read_csv(list_csv,header=None, names=['img_paths','img_labels'])
        self.root_path = root_path
        self.img_paths = listpd.img_paths.tolist()
        self.img_labels = listpd.img_labels.tolist()
        self.transform = transform

    def __len__(self):
        return len(self.img_paths)

    def __getitem__(self, item):
        img_path = os.path.join(self.root_path, self.img_paths[item])
        label = self.img_labels[item]
        img = cv2.imread(img_path)  # [h,w,3], BGR
        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)  # [132,106]
        if self.transform is not None:
            img = self.transform(img)

        return torch.from_numpy(img[np.newaxis,:]).float(),label

if __name__ == '__main__':
    # data1 = MS1M(root_path='/media/hszc/data1/face_data/ms1m/ms1m_aligned',
    #              list_csv='/home/hszc/zhangchi/channel-prune/train_list.csv',
    #              transform=FaceAug(tg_size=(112,96)))
    # img = data1[10]
    # cv2.imshow('aa',img)
    # cv2.waitKey(0)

    data1 = LFW_zip( lfw_path='/media/hszc/data/lfw.zip',
                     pairtxt='/home/hszc/zhangchi/sphereface_pytorch/data/pairs.txt',
                     landmarktxt='/home/hszc/zhangchi/sphereface_pytorch/data/lfw_landmark.txt')
    print data1[0][0].size(),data1[0][2]